
import math
import numpy as np
import random


def salt_pepper_noise(image, ratio):
    output = np.zeros(image.shape, np.uint8)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):

            rand = random.random()
            if rand < ratio:  # salt pepper noise
                if random.random() > 0.5:  # change the pixel to 255
                    output[i][j] = 255
                else:
                    output[i][j] = 0
            else:
                output[i][j] = image[i][j]

    return output


def gaussian_noise_kernel(x, mu, sigma):
	return np.exp(-1 * ((x - mu) ** 2) / (2 * (sigma ** 2))) / (math.sqrt(2 * np.pi) * sigma)


def gaussian_noise(image, ratio, sigma):
    # generate gaussian kernel
    x = np.linspace(- 4 * sigma, 4 * sigma, 100)
    kernel = gaussian_noise_kernel(x, 0, sigma)

    # output image
    output = np.zeros(image.shape, np.uint8)

    for i in range(image.shape[0]):
        for j in range(image.shape[1]):

            rand = random.random()
            if rand < ratio:  # apply gaussian noise
                pos = int(100 * random.random())

                if x[pos] < 0:
                    temp = image[i][j] * (1 - kernel[pos])
                    if temp < 0:
                        output[i][j] = 0
                    else:
                        output[i][j] = temp
                    continue

                if x[pos] >= 0:
                    temp = image[i][j] * (1 + kernel[pos])
                    if temp > 255:
                        output[i][j] = 255
                    else:
                        output[i][j] = temp

            else:
                output[i][j] = image[i][j]

    return output
